import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd.function import Function
import math
import re
import numpy as np
import logging

def replace_prune_module(mod: nn.Module, mapping: dict, pattern='.*', prefix=None, device=None):
    if device is None:
        device = list(mod.parameters())[0].device

    for name, m in mod.named_children():
        fullname = prefix + '.' + name if prefix else name
        if type(m) in mapping and re.search(pattern, fullname):
            dst_type = mapping[type(m)]
            logging.info(f'{fullname}: {type(m).__name__} is replaced with {dst_type.__name__}')
            mod._modules[name] = dst_type.from_dense(m).to(device)
        else:
            replace_prune_module(m, mapping, pattern=pattern, prefix=fullname, device=device)